#Anaconda/envs/pysagin python
# -*- coding: utf-8 -*-
'''
airbase.py
=========

.. module:: airbase
  :platform: Windows, Linux
  :synopsis: 空基平台

.. moduleauthor:: WangXi

简介
----

版本
----

- 版本 1.0 (2025/03/28): 初始版本

'''

from astropy.time import Time
import astropy.units as u
import numpy as np

from pysagin.base import Base

class AirBase(Base):
    def __init__(self, name:str, **kwargs):
        super().__init__(name, **kwargs)


class Trajectory:
    def __init__(self, name:str, avg_speed:float, **kwargs):
        """
        飞行器航机规划

        Args:
            name (str): 航迹名称
            avg_speed (float): 平均航速 (m/s)
        """
        self.name = name
        self.avg_speed = avg_speed  # 平均航速 (m/s)
        self.waypoints = [] # 航点列表 [(纬度, 经度, 高度, 累计时间)]

    def add_straight(self, start:tuple, end:tuple):
        """添加两点之间的直线飞行航迹

        根据起始点和终止点坐标,计算无人机的直线飞行路径,
        航点数量和飞行时间由平均速度自动确定,距离计算考虑了地球曲率近似。

        Args:
            start (tuple): 起点坐标,格式为 (经度, 纬度, 高度)
                        示例: (120.0, 30.0, 0) 表示东经120度北纬30度,海拔0米
            end (tuple): 终点坐标,格式为 (经度, 纬度, 高度)
                        示例: (120.1, 30.1, 100) 表示东经120.1度北纬30.1度,海拔100米

        Returns:
            None: 航点会添加到self.waypoints列表中,格式为 (纬度, 经度, 高度, 累计时间)

        Notes:
            1. 距离计算采用近似方法: 111320米≈1纬度,经度距离按cos(纬度)缩放
            2. 高度差直接参与三维距离计算
            3. 默认生成10个均匀分布的航点
            4. 时间累计计算确保多段航迹连续

        Example:
            >>> 轨迹生成器 = TrajectoryGenerator(avg_speed=10.0)  # 10米/秒
            >>> 轨迹生成器.add_straight((120.0, 30.0, 0), (120.0, 30.1, 100))
            # 生成从30°N到30.1°N的直线路径,经度保持120°E,爬升至100米高度
            # 飞行时间约为1113.2秒(距离11132米/速度10m/s)
        """
        # 解构起点和终点坐标
        lon1, lat1, alt1 = start  # 起点经度/纬度/高度
        lon2, lat2, alt2 = end    # 终点经度/纬度/高度
        
        # 计算三维距离(考虑地球曲率近似)
        # 纬度距离: 1度 ≈ 111320米
        lat_dist = (lat2 - lat1) * 111320  
        # 经度距离: 1度 ≈ 111320*cos(纬度)米 
        lon_dist = (lon2 - lon1) * 111320 * np.cos(np.radians((lat1 + lat2)/2))
        # 高度差直接取差值
        alt_dist = alt2 - alt1  
        
        # 计算总距离和所需时间
        dist = np.sqrt(lat_dist**2 + lon_dist**2 + alt_dist**2)  # 三维欧氏距离
        duration = dist / self.avg_speed  # 飞行时间(秒)

        # 生成插值点(默认10个航点)
        t = np.linspace(0, duration, 10)  # 时间轴均匀分布
        lat = np.linspace(lat1, lat2, len(t))  # 纬度线性插值
        lon = np.linspace(lon1, lon2, len(t))  # 经度线性插值
        alt = np.linspace(alt1, alt2, len(t))  # 高度线性插值
        
        # 计算时间基准(累计之前所有航段时间)
        base_time = self.waypoints[-1][3] if self.waypoints else 0
        
        # 添加航点到列表
        for ti, la, lo, al in zip(t, lat, lon, alt):
            self.waypoints.append((
                la,      # 纬度(度)
                lo,      # 经度(度) 
                al,      # 高度(米)
                ti + base_time  # 累计时间(秒)
            ))
    
    def add_circle(self, center, radius, loops=1):
        """添加圆形飞行轨迹
        
        根据中心点坐标和半径生成圆形飞行路径,支持多圈绕飞,
        自动计算飞行时间并保持恒定速度。

        Args:
            center (tuple): 圆心坐标 (经度, 纬度, 高度)
                            示例: (120.0, 30.0, 100) 表示北纬30度东经120度,海拔100米
            radius (float): 圆半径(单位:米)
                            示例: 50 表示50米半径的圆形
            loops (int):    绕飞圈数,默认为1
                            示例: 2 表示绕飞两圈

        Returns:
            None: 航点会添加到self.waypoints列表中,格式为 (纬度, 经度, 高度, 累计时间)

        Notes:
            1. 通过极坐标方程生成圆形路径点
            2. 经纬度转换考虑地球曲率(111320米/度)
            3. 经度距离按cos(纬度)缩放
            4. 自动计算飞行时间保持恒定速度
            纬度 = 中心纬度 + (radius/111320) * cos(θ)
            经度 = 中心经度 + (radius/(111320*cos(纬度))) * sin(θ)
            其中θ为极角

        Example:
            >>> 轨迹生成器 = TrajectoryGenerator(avg_speed=5.0)  # 5米/秒
            >>> 轨迹生成器.add_circle((30.0, 120.0, 100), 50, 2)
            # 生成100米高度,半径50米的圆形轨迹,绕飞2圈
            # 总距离 = 2 * π * 50 * 2 ≈ 628.3米
            # 飞行时间 ≈ 628.3/5 ≈ 125.66秒
        """
        # 解构中心点坐标
        lon0, lat0, alt0 = center  # 圆心经度/纬度/高度
        
        # 计算圆形周长和总飞行时间
        circumference = 2 * np.pi * radius  # 单圈周长
        duration = circumference * loops / self.avg_speed  # 总飞行时间
        
        # 生成极角参数(支持多圈绕飞)
        theta = np.linspace(0, 2*np.pi*loops, 50)  # 默认50个航点/圈
        t = np.linspace(0, duration, len(theta))   # 时间轴
        
        # 地球曲率常数(米/度)
        meters_per_degree = 111320
        
        # 生成圆形航点
        for angle, time in zip(theta, t):
            # 纬度计算(直接使用弧度值)
            lat = lat0 + (radius / meters_per_degree) * np.cos(angle)
            
            # 经度计算(考虑纬度缩放)
            lon = lon0 + (radius / (meters_per_degree * np.cos(np.radians(lat0)))) * np.sin(angle)
            
            # 添加航点(时间累计计算)
            self.waypoints.append((
                lat,  # 纬度(度)
                lon,  # 经度(度)
                alt0,  # 保持固定高度(米)
                time + (self.waypoints[-1][3] if self.waypoints else 0)  # 累计时间(秒)
            ))
            
    def add_figure8(self, center, size, loops=1):
        """添加八字飞行轨迹(双纽线)
        
        生成标准的八字形飞行路径(数学上称为双纽线),支持多圈绕飞,
        自动根据无人机平均速度计算飞行时间。

        Args:
            center (tuple): 轨迹中心点 (经度, 纬度, 高度)
                        示例: (120.0, 30.0, 100) 表示北纬30度东经120度,海拔100米
            size (float):   轨迹尺寸(单位:米),控制整体大小
                        示例: 50 表示宽度约50米的八字轨迹
            loops (int):    绕行圈数,默认为1
                        示例: 2 表示绕飞两个完整的八字

        Returns:
            None: 航点会添加到self.waypoints列表中,格式为 (纬度, 经度, 高度, 累计时间)

        Notes:
            1. 使用双纽线参数方程:
            x = a * sinθ
            y = a * sinθ * cosθ
            2. 周长近似公式:L ≈ 5.244 * size
            3. 经纬度转换:
            纬度偏移 = (a * sinθ) / 111320
            经度偏移 = (a * sinθ * cosθ) / (111320 * cos(纬度))

            1. 每圈默认生成50个航点
            2. 自动计算总飞行时间(总距离/平均速度)
            3. 经度距离补偿了纬度缩放因子

        Example:
            >>> 轨迹生成器 = TrajectoryGenerator(avg_speed=2.0)  # 2米/秒
            >>> 轨迹生成器.add_figure8((120.0, 30.0, 50), 30, 3)
            # 生成高度50米,尺寸30米的八字轨迹,绕飞3圈
            # 总距离 ≈ 5.244 * 30 * 3 ≈ 472米
            # 飞行时间 ≈ 472/2 ≈ 236秒
        """
        # 解构中心点坐标
        lon0, lat0, alt0 = center
        
        # 1. 计算轨迹总长度和飞行时间
        single_loop_length = 5.244 * size  # 双纽线周长近似值
        total_length = single_loop_length * loops
        duration = total_length / self.avg_speed
        
        # 2. 生成参数化角度(支持多圈)
        theta = np.linspace(0, 2 * np.pi * loops, 50 * loops)  # 50点/圈
        t = np.linspace(0, duration, len(theta))  # 时间轴
        
        # 3. 双纽线参数计算
        a = size / 111320  # 将米转换为纬度偏移量
        
        # 4. 生成航点
        base_time = self.waypoints[-1][3] if self.waypoints else 0  # 获取当前最后时间
        for angle, time in zip(theta, t):
            # 纬度计算(正弦分量)
            lat = lat0 + a * np.sin(angle)
            
            # 经度计算(正弦*余弦分量,考虑纬度缩放)
            lon = lon0 + (a * np.sin(angle) * np.cos(angle)) / np.cos(np.radians(lat0))
            
            # 添加航点(固定高度,累计时间)
            self.waypoints.append((
                lat,    # 纬度(度)
                lon,    # 经度(度)
                alt0,   # 高度(米)
                base_time + time  # 累计时间(秒)
            ))
    
    def query_position(self, time):
        """查询无人机在指定时刻的预期位置
        
        通过线性插值计算任意时刻的经纬度和高度，
        当查询时间超过总时长时返回最后一个航点位置。

        Args:
            time (float): 查询时间（单位：秒）
                        示例: 25.3 表示飞行开始后25.3秒时的位置

        Returns:
            tuple: 三维坐标 (纬度, 经度, 高度)
                纬度/经度单位：度
                高度单位：米

        Notes:
            1. 提取所有航点的时间序列和坐标序列
            2. 使用numpy的interp函数进行线性插值
            3. 对纬度、经度、高度分别进行插值计算
            4. 超时情况返回最后一个航点

            边界处理:
                - 时间<0: 返回起点坐标
                - 时间>总时长: 返回终点坐标
                - 空航点列表: 返回(0,0,0)

            计算复杂度:
                O(n) 其中n为航点数量

        Example:
            >>> 轨迹生成器 = Trajectory(avg_speed=5.0)
            >>> 轨迹生成器.add_straight((30.0,120.0,0), (30.1,120.0,100))
            >>> 轨迹生成器.query_position(10.0)
            (30.05, 120.0, 50.0)  # 飞行中途点
        """
        # 边界条件检查
        if not self.waypoints:
            return (0.0, 0.0, 0.0)  # 空航点列表返回原点
        
        # 提取时间序列和各坐标轴数据
        times = np.array([wp[3] for wp in self.waypoints])  # 所有航点的时间戳
        lats = np.array([wp[0] for wp in self.waypoints])   # 纬度序列
        lons = np.array([wp[1] for wp in self.waypoints])   # 经度序列
        alts = np.array([wp[2] for wp in self.waypoints])   # 高度序列
        
        # 处理超时情况
        if time <= times[0]:
            return self.waypoints[0][:3]  # 返回起点
        if time >= times[-1]:
            return self.waypoints[-1][:3]  # 返回终点
        
        # 三维线性插值
        lat = np.interp(time, times, lats)  # 纬度插值
        lon = np.interp(time, times, lons)  # 经度插值
        alt = np.interp(time, times, alts)  # 高度插值
        
        return (lat, lon, alt)
    
   

class UavBase(AirBase):
    def __init__(self, name: str, init_time: Time, init_gps: list, **kwargs):
        """
        初始化无人机对象。

        Args:
            name (str): 无人机名称
            init_time (astropy.time.Time): 初始时间
            init_gps (list): 初始GPS坐标 [经度, 纬度, 高度]
        """
        super().__init__(name, **kwargs)
        self.init_time = init_time
        self.current_time = init_time
        self.init_gps = init_gps
        self.current_gps = init_gps

        # 航线相关参数
        self.speed = None           # 航速, 单位km/s
        self.departure_time = None  # 起飞时间
        self.landing_time = None    # 降落时间
        self.trajectory:Trajectory = None      # 航线轨迹

    def set_trajectory(self, trajectory:Trajectory, departure_time:Time) -> None:
        """
        规划航线, 并设置起飞时间。

        Args:
        
        Note:
            飞行器降落时间根据最后一个航线点设置获取
            
        """
        self.trajectory = trajectory
        self.departure_time = departure_time
        self.landing_time = self.departure_time + self.trajectory.waypoints[-1][3] * u.s

    def update(self, time: Time) -> None:
        """
        更新无人机的位置。

        - 如果未规划航线，保持在初始位置。
        - 如果当前时间小于出发时间，返回出发点。
        - 如果当前时间大于等于出发时间且未到达目的地，计算当前位置。
        - 如果已经到达目的地，返回目的地坐标。
        """
        self.current_time = time

        # 检查是否已规划航线
        if self.trajectory is None:
            raise ValueError("航线未规划，请先调用 set_trajectory 设置航线")

        # 当前时间早于出发时间，返回出发点
        if time < self.departure_time:
            self.current_gps = [self.trajectory.waypoints[0][1],self.trajectory.waypoints[0][0],self.trajectory.waypoints[0][2]]
            return

        # 如果飞行时间超过最大飞行时间，返回目的地
        if time >= self.landing_time:
            self.current_gps = [self.trajectory.waypoints[-1][1],self.trajectory.waypoints[-1][0],self.trajectory.waypoints[-1][2]]
            return
        
        delta_time_s = (time - self.departure_time).to(u.s).value  # 当前时间与起飞时间差值(秒)
        current_lat, current_lon, current_alt = self.trajectory.query_position(delta_time_s)  # 查询当前位置

        self.current_gps = [current_lon, current_lat, current_alt/1000.0]


if __name__ == '__main__':
    # 创建无人机对象
    uav = UavBase(
        name="UAV-001",
        init_time=Time("2025-03-28 00:00:00"),
        init_gps=[120.0, 30.0, 0.0]
    )

    uav_trajectory = Trajectory(name="UAV-001",avg_speed=5.0)
    uav_trajectory.add_straight(start=(30.0, 120.0, 0), end=(30.1, 120.1, 100))
    uav_trajectory.add_figure8(center=(30.1, 120.1, 100), size=1000, loops=3)
    uav_trajectory.add_straight(start=(30.1, 120.1, 100), end=(30.2, 120.2, 80))
    uav_trajectory.add_circle(center=(30.2, 120.2, 80), radius=1500, loops=2)
    uav_trajectory.add_straight(start=(30.2, 120.2, 80), end=(30.0, 120.0, 0))

    uav.set_trajectory(trajectory=uav_trajectory, departure_time=Time("2025-03-28 00:00:00"))

    # 更新位置
    import matplotlib.pyplot as plt
    
    uav.update(Time("2025-03-27 12:05:00"))
    [lon, lat, alt] = uav.current_gps
    plt.scatter(lon, lat, c='red', s=20)
    print("当前位置:", uav.current_gps)
    
    uav.update(Time("2025-03-28 01:20:00"))
    [lon, lat, alt] = uav.current_gps
    plt.scatter(lon, lat, c='blue', s=20)
    print("当前位置:", uav.current_gps)
    
    uav.update(Time("2025-03-28 10:00:00"))
    [lon, lat, alt] = uav.current_gps
    plt.scatter(lon, lat, c='blue', s=20)
    print("当前位置:", uav.current_gps)
    
    uav.update(Time("2025-03-28 20:00:00"))
    [lon, lat, alt] = uav.current_gps
    plt.scatter(lon, lat, c='blue', s=20)
    print("当前位置:", uav.current_gps)

    uav.update(Time("2025-03-29 14:00:00"))
    [lon, lat, alt] = uav.current_gps
    plt.scatter(lon, lat, c='green', s=20)
    print("当前位置:", uav.current_gps)
    
    lons = [wp[1] for wp in uav.trajectory.waypoints]
    lats = [wp[0] for wp in uav.trajectory.waypoints]
    plt.plot(lons, lats, 'b-')
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")
    plt.show()